/*
 * JGO2RDF creates a RDF representation of GAF2 files, which are automatically sourced from GO ftp site.
 * Copyright (C) 2012 Matthew M. Hindle & Artem Lysenko
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package uk.co.jgo2rdf.goa.go;

import org.apache.log4j.Logger;

import uk.co.jgo2rdf.goa.utils.DefaultSuppliers;

import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.google.common.collect.SetMultimap;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.*;

/**
 * @author hindlem
 *         Created 05-May-2010 13:42:20
 */
public class GOIndexer implements GoIndex  {

    private static final Logger log = Logger.getLogger(GOIndexer.class);

    public static String PROC_ROOT = "GO:0008150".intern();
    public static String CellComp_ROOT = "GO:0005575".intern();
    public static String Func_ROOT = "GO:0003674".intern();

    private Set<String> blacklistTerms = new HashSet<String>();
    private Map<String, Map<String, String>> ontology2Name2term = new HashMap<String, Map<String, String>>();

    private Map<String, String> goId2Name = new HashMap<String, String>();
    private Map<String, String> goId2Description = new HashMap<String, String>();

    private SetMultimap<String, String> goFromTo =  Multimaps.newSetMultimap(Maps.<String, Collection<String>> newHashMap(), DefaultSuppliers.set(String.class));
    private Map<String, OntologyType> goToOntology = new HashMap<String, OntologyType>();
    private SetMultimap<String, String> ontologyToGo =  Multimaps.newSetMultimap(Maps.<String, Collection<String>> newHashMap(), DefaultSuppliers.set(String.class));
    private Map<String, String> ontologyToRoot = new HashMap<String, String>();
    private Map<String, String> alt_id_2_realId = new HashMap<String, String>();
    
    private SetMultimap<String, String> ancestorIndex =  Multimaps.newSetMultimap(Maps.<String, Collection<String>> newHashMap(), DefaultSuppliers.set(String.class));
    private SetMultimap<String, String> childrenIndex =  Multimaps.newSetMultimap(Maps.<String, Collection<String>> newHashMap(), DefaultSuppliers.set(String.class));


    public GOIndexer(InputStream is) throws IOException, OntologyType.UnknownOntologyNameException {
        log.info("initializing GO Indexer");
        index(is);
    }

    private void index(InputStream is) throws IOException, OntologyType.UnknownOntologyNameException {
        BufferedReader br = new BufferedReader(new InputStreamReader(is));
        try {
            String term = null;
            String name = null;
            String ontology = null;
            String definition = null;
            Set<String> synonyms = new HashSet<String>();
            Set<String> altIds = new HashSet<String>();

            while (br.ready()) {
                String line = br.readLine();

                if (line.trim().equalsIgnoreCase("[Term]")) {
                    fillNames(term, ontology, name, definition, synonyms);
                    ontology = null;
                    name = null;
                    definition = null;
                    altIds.clear();
                    synonyms.clear();
                    term = br.readLine().substring(4).trim().toUpperCase().intern();
                }
                else if(term != null){
                	if (line.startsWith("name:")) {
                        name = line.substring(5).replaceAll("\"", "").trim();
                        goId2Name.put(term, line.substring(5).replaceAll("\"", "").trim());
                    } else if (line.startsWith("synonym:") && term != null) {
                        int from = line.indexOf("\"");
                        int to = line.indexOf("\"", from + 1);
                        synonyms.add(line.substring(from, to).replaceAll("\"", "").trim().toUpperCase());
                    } else if (line.startsWith("def:")) {
                        int from = line.indexOf("\"");
                        int to = line.indexOf("\"", from + 1);
                        definition = line.substring(from, to).replaceAll("\"", "").trim().toUpperCase();
                        String raw_def = line.substring(from, to).replaceAll("\"", "").trim();
                        goId2Description.put(term, raw_def);
                        for (String id : altIds)
                            goId2Description.put(id, raw_def);
                    } else if (line.startsWith("is_a:")) {
                        String termTo = line.substring(5, line.indexOf('!') - 1).trim().toUpperCase().intern();
                        goFromTo.put(term, termTo);
                    } else if (line.startsWith("relationship: part_of")) {
                        String termTo = line.substring("relationship: part_of".length(), line.indexOf('!') - 1).trim().toUpperCase().intern();
                        goFromTo.put(term, termTo);
                    } else if (line.startsWith("namespace:")) {
                        ontology = line.substring("namespace:".length()).trim().toUpperCase().intern();
                        goToOntology.put(term, OntologyType.getOntology(ontology));
                        if (name.toUpperCase().equalsIgnoreCase(ontology)) {
                            ontologyToRoot.put(ontology, term);
                        }
                        ontologyToGo.put(ontology, term);
                    } else if (line.startsWith("alt_id")) {
                        String termTo = line.substring("alt_id:".length()).trim().toUpperCase().intern();
                        if (alt_id_2_realId.put(termTo, term) != null) {
                            System.err.println("alt_id " + termTo + " is assigned to multiple GO Terms");
                        }
                        altIds.add(termTo);
                        goId2Name.put(termTo, name);
                        goToOntology.put(termTo, OntologyType.getOntology(ontology));
                    }
                }	
            }
            fillNames(term, ontology, name, definition, synonyms);
            synonyms.clear();

            log.info("finished processing GO file, continuing to build DAG index. " + goToOntology.size() + " terms indexed so far");

            for (String goTerm : goFromTo.keySet()) {
                Set<String> terms = new HashSet<String>();
                getAncestors(goTerm, terms);
                //System.out.println(goTerm+"-- "+Arrays.toString(terms.toArray()));
                terms.remove(goTerm);
                //System.out.println(goTerm+"++ "+Arrays.toString(terms.toArray()));
                ancestorIndex.putAll(goTerm, terms);
            }

            for (Map.Entry<String, String> goTerm : ancestorIndex.entries()) {
            	childrenIndex.put(goTerm.getValue(), goTerm.getKey());
            }
        } finally {
            br.close();
        }
        log.info("GO index complete");
    }

    private void fillNames(String term, String ontology, String name, String definition, Set<String> synonyms) {
        if (term != null && ontology != null) {
            Map<String, String> name2term = ontology2Name2term.get(ontology);
            if (name2term == null) {
                name2term = new HashMap<String, String>();
                ontology2Name2term.put(ontology, name2term);
            }
            if (name != null && null != name2term.put(name, term)) {
                blacklistTerms.add(name);
            }

            if (definition != null && null != name2term.put(definition, term)) {
                blacklistTerms.add(definition);
            }

            for (String synonym : synonyms) {
                if (null != name2term.put(synonym, term)) {
                    blacklistTerms.add(synonym);
                }
            }
        }

    }

    public Set<String> getAncestors(String goTerm) {
        Set<String> a = ancestorIndex.get(goTerm);
        HashSet<String> ancestors = new HashSet<String>();
        if (a != null) ancestors.addAll(a);
        return ancestors;
    }

    @Override
    public Set<String> getChildren(String goTerm) {
        Set<String> c = childrenIndex.get(goTerm);
        HashSet<String> children = new HashSet<String>();
        if (c != null) children.addAll(c);
        return children;
    }

    /**
     * is goTerm1 and ancestor of goTerm2
     *
     * @param goTerm1
     * @param goTerm2
     * @return
     */
    @Override
    public boolean isAncestor(String goTerm1, String goTerm2) {
        Set<String> ancestors = ancestorIndex.get(goTerm2);
        return !(ancestors == null || ancestors.size() == 0) && ancestors.contains(goTerm1);
    }

    public Set<String> getAncestors(Set<String> goTerms) {
        Set<String> terms = new HashSet<String>();
        for (String goTerm : goTerms) {
            Set<String> ancestors = getAncestors(goTerm);
            if (ancestors != null)
                terms.addAll(ancestors);
        }
        return terms;
    }

    public Set<String> getImmediateAncestors(String goTerm) {
        Set<String> terms = new HashSet<String>();

        Set<String> ancestorsPartOf = goFromTo.get(goTerm);
        if (ancestorsPartOf != null)
            terms.addAll(ancestorsPartOf);
        return terms;
    }

    private void getAncestors(String goTerm, Set<String> ancestorCollector) {
        Set<String> ancestors = goFromTo.get(goTerm);

        if (ancestors != null) {
            ancestorCollector.addAll(ancestors);
            for (String ancestor : ancestors) {
                getAncestors(ancestor, ancestorCollector);
            }
        }
    }

    public String getName(String GOid) {
        return goId2Name.get(GOid);
    }

    public String getDescription(String GOid) {
        return goId2Description.get(GOid);
    }

    public String getGOID(String ontology, String name) {
        if (blacklistTerms.contains(name)) {
            System.err.println("term name is ambiguous " + name);
            return null;
        }

        Map<String, String> nameMap = ontology2Name2term.get(ontology);
        if (nameMap != null) {
            String result = nameMap.get(name.toUpperCase());
            if (result == null) {
                result = nameMap.get((name + " activity").toUpperCase());
            }
            if (result == null) {
                result = nameMap.get((name.replaceAll(" ", ",") + " activity").toUpperCase());
            }
            return result;
        }
        throw new RuntimeException("Ontology is invalid: " + ontology + " (use, " + Arrays.toString(ontology2Name2term.keySet().toArray(new String[ontology2Name2term.keySet().size()])) + ")");
    }

    public Set<String> upgradeObsoleteTerms(Collection<String> ids) {
        String[] newids = ids.toArray(new String[ids.size()]);
        for (int i = 0; i < newids.length; i++) {
            String newId = alt_id_2_realId.get(newids[i]);
            if (newId != null) {
                newids[i] = newId;
            }
        }
        return new HashSet(Arrays.asList(newids));
    }


    public Set<String> removeRedundentTerms(Collection<String> ids) {
        Set<String> newIds = new HashSet<String>(ids);
        for (String id : ids) {
            Set<String> ancestors = ancestorIndex.get(id);
            if (ancestors != null)
                newIds.removeAll(ancestors);
        }
        if (newIds.size() == 0) {
            for (String id : ids)
                System.err.println(id + ":" + Arrays.toString(ancestorIndex.get(id).toArray()));
            throw new RuntimeException(Arrays.toString(ids.toArray()));
        }
        return newIds;
    }

    /**
     * @param ids   to filter aginst toSet
     * @param toSet checks for redundancy to here
     * @return
     */
    public Set<String> removeRedundentTerms(Collection<String> ids, Collection<String> toSet) {
        Set<String> newIds = new HashSet<String>(ids);
        newIds.removeAll(toSet);
        for (String id : toSet) {
            Set<String> ancestors = ancestorIndex.get(id);
            if (ancestors != null)
                newIds.removeAll(ancestors);
        }
        return newIds;
    }

    /**
     * @param ids   to filter aginst toSet
     * @param toSet checks for non consensus to here
     * @return
     */
    public Set<String> removeConsensualTerms(Collection<String> ids, Collection<String> toSet) {
        Set<String> newIds = removeRedundentTerms(ids, toSet);
        for (String id : ids) {
            Set<String> ancestors = ancestorIndex.get(id);
            if (ancestors != null) {
                for (String ancestor : ancestors) {
                    if (toSet.contains(ancestor)) {
                        newIds.remove(id);
                    }
                }
            }
        }
        return newIds;
    }

    /**
     * @param key
     * @param value
     * @param map   map to add to
     * @param <K>   key type
     * @param <V>   value type
     */
    public <K, V> void getAndUpdate(K key, V value, Map<K, Set<V>> map) {
        Set<V> valueSet = map.get(key);
        if (valueSet == null) {
            valueSet = new HashSet<V>();
            map.put(key, valueSet);
        }
        valueSet.add(value);
    }


    public OntologyType getOntology(String goid) {
        return goToOntology.get(goid);
    }

    @Override
    public Set<String> getOntologies() {
        return ontologyToGo.keySet();
    }

    @Override
    public Set<String> getTermsInOntology(String ontology) {
        return ontologyToGo.get(ontology);
    }

    @Override
    public String getRootTermInOntology(String ontology) {
        return ontologyToRoot.get(ontology);  //To change body of implemented methods use File | Settings | File Templates.
    }

    public String getRealId(String string) {
        return alt_id_2_realId.get(string);
    }

    public Set<String> removeObsoleteTerms(Set<String> goTerms) {
        Set<String> newSet = new HashSet<String>();
        for (String term : goTerms) {
            if (goToOntology.containsKey(term))
                newSet.add(term);
        }
        return newSet;
    }
}